# coding: utf-8
import nltk
import os
import codecs
import sys
import random
import sys

filename = sys.argv[1]
raw_input('{}? enter to continue, ctrl+c to exit: '.format(filename))
train_size = 0.8
valid_size = 0.1
test_size = 0.1

assert train_size + valid_size + test_size == 1.0
dirname = os.path.dirname(filename)
train_filename = os.path.join(dirname, 'train.txt')
valid_filename = os.path.join(dirname, 'valid.txt')
test_filename = os.path.join(dirname, 'test.txt')

all_lines = []
with codecs.open(filename, 'r', encoding='utf-8') as f:
    for line in f:
        all_lines.append(line.strip())

random.shuffle(all_lines)

l = len(all_lines)
sp1 = int(l * (train_size))
sp2 = int(l * (train_size + valid_size))

with codecs.open(train_filename, 'w', encoding='utf-8') as f:
    for line in all_lines[:sp1]:
        f.write(line)
        f.write('\n')
    print('{} ok!'.format(train_filename))

with codecs.open(valid_filename, 'w', encoding='utf-8') as f:
    for line in all_lines[sp1:sp2]:
        f.write(line)
        f.write('\n')
    print('{} ok!'.format(valid_filename))

with codecs.open(test_filename, 'w', encoding='utf-8') as f:
    for line in all_lines[sp1:sp2]:
        f.write(line)
        f.write('\n')
    print('{} ok!'.format(test_filename))
